// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2015 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_TYPE_CASTING_AVX_H
#define EIGEN_TYPE_CASTING_AVX_H

namespace Eigen {

namespace internal {

    // For now we use SSE to handle integers, so we can't use AVX instructions to cast
    // from int to float
    template <> struct type_casting_traits<float, int>
    {
        enum
        {
            VectorizedCast = 0,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

    template <> struct type_casting_traits<int, float>
    {
        enum
        {
            VectorizedCast = 0,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

#ifndef EIGEN_VECTORIZE_AVX512

    template <> struct type_casting_traits<Eigen::half, float>
    {
        enum
        {
            VectorizedCast = 1,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

    template <> struct type_casting_traits<float, Eigen::half>
    {
        enum
        {
            VectorizedCast = 1,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

    template <> struct type_casting_traits<bfloat16, float>
    {
        enum
        {
            VectorizedCast = 1,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

    template <> struct type_casting_traits<float, bfloat16>
    {
        enum
        {
            VectorizedCast = 1,
            SrcCoeffRatio = 1,
            TgtCoeffRatio = 1
        };
    };

#endif  // EIGEN_VECTORIZE_AVX512

    template <> EIGEN_STRONG_INLINE Packet8i pcast<Packet8f, Packet8i>(const Packet8f& a) { return _mm256_cvttps_epi32(a); }

    template <> EIGEN_STRONG_INLINE Packet8f pcast<Packet8i, Packet8f>(const Packet8i& a) { return _mm256_cvtepi32_ps(a); }

    template <> EIGEN_STRONG_INLINE Packet8i preinterpret<Packet8i, Packet8f>(const Packet8f& a) { return _mm256_castps_si256(a); }

    template <> EIGEN_STRONG_INLINE Packet8f preinterpret<Packet8f, Packet8i>(const Packet8i& a) { return _mm256_castsi256_ps(a); }

    template <> EIGEN_STRONG_INLINE Packet8f pcast<Packet8h, Packet8f>(const Packet8h& a) { return half2float(a); }

    template <> EIGEN_STRONG_INLINE Packet8f pcast<Packet8bf, Packet8f>(const Packet8bf& a) { return Bf16ToF32(a); }

    template <> EIGEN_STRONG_INLINE Packet8h pcast<Packet8f, Packet8h>(const Packet8f& a) { return float2half(a); }

    template <> EIGEN_STRONG_INLINE Packet8bf pcast<Packet8f, Packet8bf>(const Packet8f& a) { return F32ToBf16(a); }

}  // end namespace internal

}  // end namespace Eigen

#endif  // EIGEN_TYPE_CASTING_AVX_H
